package com.ai.demo.controller;

import com.ai.demo.entity.vo.Result;
import com.ai.demo.enums.ChatTypeEnum;
import com.ai.demo.service.IFileService;
import com.ai.demo.service.ISpringAiChatRecordService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.core.io.Resource;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.Objects;

@Slf4j
@RequiredArgsConstructor
@RestController
@RequestMapping("/ai/pdf")
public class PdfController {

    private final IFileService fileService;

    private final ChatClient pdfChatClient;

    private final ISpringAiChatRecordService springAiChatRecordService;

    /**
     * 上传文件
     * @param chatId
     * @param file
     * @return
     */
    @RequestMapping("/upload/{chatId}")
    public Result upload(@PathVariable String chatId, @RequestParam("file") MultipartFile file){
        try {
            //1.校验Pdf格式文件
            if (!Objects.equals(file.getContentType(), "application/pdf")) {
                return Result.fail("只能上传PDF文件！");
            }
            //2.保存文件
            boolean flag = fileService.save(chatId, file.getResource());
            if (!flag) {
                return Result.fail("保存文件失败！");
            }
            return Result.ok();
        } catch (Exception e) {
            log.error("上传文件失败！", e);
            return Result.fail("上传文件失败！");
        }
    }

    /**
     * 下载文件
     * @param chatId
     */
    @GetMapping("/file/{chatId}")
    public ResponseEntity<Resource> download(@PathVariable("chatId") String chatId){
        //1.读取文件
        Resource resource = fileService.getFile(chatId);
        if (!resource.exists()) {
            return ResponseEntity.notFound().build();
        }
        //2.文件名编码，写入响应头
        String filename = URLEncoder.encode(Objects.requireNonNull(resource.getFilename()), StandardCharsets.UTF_8);
        //3.返回文件
        return ResponseEntity.ok()
                .contentType(MediaType.APPLICATION_OCTET_STREAM)
                .header("Content-Disposition", "attachment; filename=\"" + filename + "\"")
                .body(resource);
    }

    @RequestMapping(value = "/chat",produces = "text/html;charset=utf-8")
    public Flux<String> chat(@RequestParam String prompt, String chatId) {
        System.out.println("prompt: " + prompt);
        springAiChatRecordService.saveRecord(ChatTypeEnum.PDF.getCode(),chatId);
        return pdfChatClient.prompt()
                .user(prompt)
                .advisors(obj->obj.param(ChatMemory.CONVERSATION_ID,chatId))
                .advisors(a -> a.param(QuestionAnswerAdvisor.FILTER_EXPRESSION, "chat_id == '"+chatId+"'"))
                .stream()
                .content();
    }

}
